{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d6da08f-a7cc-4976-be15-db51e2fce32f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torchvision import transforms, datasets\n",
    "from pathlib import Path\n",
    "import mltools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd49465",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist(path=\"../data\", batch_size=100):\n",
    "    download = False if Path(f\"{path}/MNIST\").exists() else True\n",
    "    trans = transforms.ToTensor()  # 数据集格式转换\n",
    "    train_data = datasets.MNIST(root=path, train=True, transform=trans, download=download)\n",
    "    test_data = datasets.MNIST(root=path, train=False, transform=trans, download=download)\n",
    "    train_data, val_data = mltools.split_data(train_data, [9, 1])  # 训练集和验证集比例9：1\n",
    "    return mltools.iter_data([train_data, val_data, test_data], batch_size=batch_size)  # 返回数据迭代器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb4b0ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "train_iter, val_iter, test_iter = mnist(batch_size=1000)  # 获取训练集、验证集、测试集\n",
    "model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))  # 设置模型结构\n",
    "model.to(device)\n",
    "loss = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=1e-1)\n",
    "ml = mltools.MachineLearning(\"mlp\")\n",
    "ml.add_model(model)\n",
    "epoch, timer, recorder = ml.batch_create()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21c29496",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练模型\n",
    "num_epochs = epoch(20)\n",
    "animator = ml.create_animator(\n",
    "    xlabel=\"epoch\", xlim=[0, epoch.totol_epoch + 1], ylim=-0.1, legend=[\"train loss\", \"train acc\", \"val acc\"]\n",
    ")  # 创建动画器\n",
    "for current_epoch in range(1, num_epochs + 1):\n",
    "    timer.start()  # 开始计时\n",
    "\n",
    "    # 计算训练集\n",
    "    metric_train = mltools.Accumulator(3)  # 累加器：(train_loss, train_acc, train_size)\n",
    "    model.train()  # 训练模式\n",
    "    for x, y in train_iter:\n",
    "        x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "        y_train = model(x)  # 计算模型\n",
    "        train_loss = loss(y_train, y)  # 计算损失\n",
    "        train_pred = y_train.argmax(dim=1)  # 计算预测值\n",
    "        train_acc = (train_pred == y).sum()  # 计算准确数\n",
    "\n",
    "        # 梯度更新\n",
    "        optimizer.zero_grad()\n",
    "        train_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        with torch.no_grad():\n",
    "            metric_train.add(train_loss * y.numel(), train_acc, y.numel())  # 累加损失、准确数、样本数\n",
    "    recorder[0].append(metric_train[0] / metric_train[2])  # 计算平均损失\n",
    "    recorder[1].append(metric_train[1] / metric_train[2])  # 计算准确率\n",
    "\n",
    "    # 计算验证集\n",
    "    metric_val = mltools.Accumulator(2)  # 累加器：(val_loss, val_size)\n",
    "    model.eval()  # 验证模式\n",
    "    with torch.no_grad():\n",
    "        for x, y in val_iter:\n",
    "            x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "            y_val = model(x)  # 计算模型\n",
    "            val_pred = y_val.argmax(dim=1)  # 计算预测值\n",
    "            val_acc = (val_pred == y).sum()  # 计算准确数\n",
    "            metric_val.add(val_acc, y.numel())  # 累加准确数、样本数\n",
    "    recorder[2].append(metric_val[0] / metric_val[1])  # 计算准确率\n",
    "\n",
    "    timer.stop()  # 停止计时\n",
    "\n",
    "    # 打印输出值\n",
    "    ml.logger.info(f\"train loss {recorder[0][-1]:.3f}, train acc {recorder[1][-1]:.3f}, val acc {recorder[2][-1]:.3f}\")\n",
    "    ml.print_training_time_massage(timer, num_epochs, current_epoch)\n",
    "    ml.logger.info(f\"trained on {str(device)}\")\n",
    "    animator.show(recorder.data)\n",
    "else:\n",
    "    # 打印输出值\n",
    "    ml.logger.info(f\"train loss {recorder[0][-1]:.3f}, train acc {recorder[1][-1]:.3f}, val acc {recorder[2][-1]:.3f}\")\n",
    "    ml.print_training_time_massage(timer, num_epochs, current_epoch)\n",
    "    ml.logger.info(f\"trained on {str(device)}\")\n",
    "    animator.show(recorder.data)\n",
    "ml.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ed6655-9565-45c0-a9bc-5ba660ce4246",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试模型\n",
    "model.eval()\n",
    "metric = mltools.Accumulator(2)  # 累加器：(test_acc, test_size)\n",
    "with torch.no_grad():\n",
    "    for x, y in test_iter:\n",
    "        x, y = x.to(device), y.to(device)  # 转换x、y\n",
    "        y_test = model(x)  # 计算模型\n",
    "        test_pred = y_test.argmax(dim=1)  # 计算预测值\n",
    "        test_acc = (test_pred == y).sum()  # 计算准确数\n",
    "        metric.add(test_acc, y.numel())  # 累加准确数、样本数\n",
    "ml.logger.info(f\"test acc {metric[0] / metric[1]:.3f}\")  # 计算准确率并输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d98c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 预测模型\n",
    "model.eval()\n",
    "x, y = next(iter(test_iter))  # 从测试中取一个批量\n",
    "x, y = x[:10].to(device), y[:10].to(device)  # 转换x、y\n",
    "y_pred = model(x)  # 计算模型\n",
    "y_pred = y_pred.argmax(dim=1)  # 计算预测值\n",
    "labels = [f\"real:{y[index]}\\npred:{y_pred[index]}\" for index in range(y.numel())]\n",
    "mltools.images(x.squeeze(1).cpu().numpy(), labels, shape=(2, 5))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
